Don't replace unused variables with None - #2396

Some grad tests are failing in a non-trivial way. In the case of atleast_2d, test_grad is being supplied an input which is then directly returned as output (z, y = atleast_2d(x, y)). Before this change the output y would have been replaced by None and we wouldn't hit the error raised here.

Could you please check whether the error goes away if atleast_2d and all similar operations use prims.clone and not return the input directly?

I would not expect clone (which is not what PyTorch does, it really does identity.) Maybe what is missing is something in __iter__.
Note that the grad checks are still using the old logic, which has much more elaborate handling of complex outputs, so maybe moving those from running vjp to comparing to torch grad after torch.autograd.grad(thunder_function(x), x) fixes it. I would recommend looking into whether it is test specific before trying to fix. I vaguely recall a similar problem before but I don't remember the exact issue.

You are right, @t-vi , that prims.clone is not the right approach here because it does not mimic PyTorch. However I don't think this is just limited to the tests. For example,

import torch
import thunder

class MyFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: torch.Tensor) -> torch.Tensor:
        x = x.sin()
        return torch.atleast_2d(x)

    @staticmethod
    def backward(ctx, grad_output) -> torch.Tensor:
        return grad_output

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x) -> torch.Tensor:
        return MyFunction.apply(x)

model = Model().to(dtype=torch.float64)

jitted = thunder.jit(model)

y = torch.randn(10, 10, requires_grad=True, device="cuda")
z = jitted(y)

still fails with "Encountered exception ValueError: Variable t1 is being overwritten this is not allowed while tracing Model()", though this failure is unrelated to removing this helper function. Using prims.shallow_copy fixes both issues.

Btw, which __iter__ do you mean @t-vi ?

Couldn't we del _ instead when it is assigned?

import torch
import thunder

def f(x):
    y, z = torch.var_mean(x, dim=1)
    return z

jitted = thunder.jit(f)

x = torch.randn(10, 10)
z = jitted(x)

last_trace = thunder.last_traces(jitted)[-1]
var_mean_bsym = last_trace.bound_symbols[1]
print(var_mean_bsym.flat_outs)         # [None, <TensorProxy(name="t11", dtype=thunder.dtypes.float32, shape=(10,))>]

print(var_mean_bsym)                   # (_, t11) = torch.var_mean(x, 1, keepdim=False, correction=1)
                                       #   # (t10, t11) = ltorch.var_mean(x, 1, keepdim=False, correction=1)
                                       #     # (t10, t11) = prims.var_mean(x, (1,), correction=1)

The underscore doesn't appear until the printing process and is otherwise represented as None. I attempted to modify del_last_used to also add a bsym with op prims.DEL and arg None, but this meant also changing the del_printer from arg_string = ", ".join(codeutils.prettyprint(x, literals_allowed=False) for x in arg_printables) to arg_string = ", ".join(codeutils.prettyprint(x, literals_as_underscores=True) for x in arg_printables). But that gives me the howling fantods.

I think the del_last_used could check if the return is a tuple and if it is it can just iterate over that instead of using flat proxy outs. This would allow you to detect none properly.

@t-vi I'm not understanding your last comment, or at least the motivation behind it. Correct me if I'm wrong, but it seems like you are suggesting that unused TensorProxies should continue to be replaced with None in the dce pass and that del_last_used should then tidy things up. What is the advantage of replacing an unused TensorProxy with None in a bsym's outs?

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB